#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
XOS AI模型推理脚本
用于加载训练好的模型并进行SQL生成
"""

import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
import json
import logging
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class XOSSQLGenerator:
    """XOS SQL生成器"""
    
    def __init__(self, model_path: str = "./models/xos_simple_final"):
        self.model_path = model_path
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 加载模型和tokenizer
        try:
            self.tokenizer = T5Tokenizer.from_pretrained(model_path)
            self.model = T5ForConditionalGeneration.from_pretrained(model_path)
            self.model.to(self.device)
            self.model.eval()
            logger.info(f"模型加载成功: {model_path}")
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            # 回退到基础模型
            logger.info("使用基础模型...")
            self.tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
            self.model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
            self.model.to(self.device)
            self.model.eval()
    
    def generate_sql(self, question: str, context: str = "", database: str = "xos_system") -> str:
        """
        生成SQL查询
        
        Args:
            question: 自然语言问题
            context: 上下文信息
            database: 数据库名称
            
        Returns:
            生成的SQL查询
        """
        # 构建输入提示
        input_text = f"生成SQL: {question}"
        if context:
            input_text += f" 上下文: {context}"
        if database:
            input_text += f" 数据库: {database}"
        
        # 编码输入
        inputs = self.tokenizer(
            input_text,
            return_tensors="pt",
            max_length=512,
            truncation=True
        ).to(self.device)
        
        # 生成SQL
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=256,
                num_beams=4,
                temperature=0.7,
                do_sample=True,
                early_stopping=True
            )
        
        # 解码输出
        generated_sql = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return generated_sql
    
    def batch_generate(self, questions: list) -> list:
        """批量生成SQL"""
        results = []
        for question in questions:
            sql = self.generate_sql(question)
            results.append({
                "question": question,
                "generated_sql": sql
            })
        return results
    
    def interactive_mode(self):
        """交互模式"""
        logger.info("进入交互模式，输入'quit'退出")
        
        while True:
            try:
                question = input("\n请输入自然语言查询: ").strip()
                
                if question.lower() in ['quit', 'exit', '退出']:
                    break
                
                if not question:
                    continue
                
                # 生成SQL
                sql = self.generate_sql(question)
                print(f"生成的SQL: {sql}")
                
            except KeyboardInterrupt:
                print("\n退出交互模式")
                break
            except Exception as e:
                logger.error(f"生成错误: {e}")

def test_examples():
    """测试示例"""
    generator = XOSSQLGenerator()
    
    test_cases = [
        {
            "question": "查询所有活跃用户的基本信息",
            "context": "需要获取用户的姓名、邮箱和电话"
        },
        {
            "question": "获取用户张三的所有权限",
            "context": "通过用户名查询权限信息"
        },
        {
            "question": "查看库存不足的产品",
            "context": "当前库存小于最小库存的产品"
        },
        {
            "question": "统计每个客户的订单数量和总金额",
            "context": "按客户分组统计"
        },
        {
            "question": "查询最近7天的系统日志",
            "context": "按时间倒序排列"
        },
        {
            "question": "获取动态表格products的配置信息",
            "context": "包括列配置"
        }
    ]
    
    logger.info("开始测试示例...")
    
    for i, case in enumerate(test_cases, 1):
        logger.info(f"\n=== 测试用例 {i} ===")
        logger.info(f"问题: {case['question']}")
        logger.info(f"上下文: {case.get('context', '无')}")
        
        sql = generator.generate_sql(case['question'], case.get('context', ''))
        logger.info(f"生成SQL: {sql}")

def main():
    """主函数"""
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "test":
            test_examples()
        elif sys.argv[1] == "interactive":
            generator = XOSSQLGenerator()
            generator.interactive_mode()
        else:
            print("用法: python inference.py [test|interactive]")
    else:
        # 默认运行测试
        test_examples()

if __name__ == "__main__":
    main()
